import torch.nn as nn
import torch.nn.functional as F


class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        ##LeNet 结构
        self.conv1 = nn.Conv2d(3,16,5)      ##卷积层1
        self.pool1 = nn.MaxPool2d(2,2)      ##池化层1
        self.conv2 = nn.Conv2d(16,32,5)     ##卷积层2
        self.pool2 = nn.MaxPool2d(2,2)      ##池化层2
        self.fc1 = nn.Linear(17*17*32,120)  ##全连接层1
        self.fc2 = nn.Linear(120,84)        ##全连接层2
        self.fc3 = nn.Linear(84,18)         ##全连接层3

    def forward(self,x):
        x = F.relu(self.conv1(x))
        x = self.pool1(x)
        x = F.relu(self.conv2(x))
        x = self.pool2(x)
        x = x.view(-1,17*17*32)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x